import { Injectable } from '@nestjs/common'

export interface WorkflowNode {
    id: number
    type: string
    inputs: Record<string, any>
    class_type: string
}

export interface Workflow {
    nodes: WorkflowNode[]
    edges: Array<{
        from_node: number
        from_output: string
        to_node: number
        to_input: string
    }>
}

export interface SamplerConfig {
    name: string
    scheduler: string
    steps: number
    cfg: number
    denoise: number
    seed: number
}

export const SAMPLER_TYPES = {
    EULER: 'euler',
    EULER_A: 'euler_ancestral',
    HEUN: 'heun',
    DPM_2: 'dpm_2',
    DPM_2_A: 'dpm_2_ancestral',
    LMS: 'lms',
    DDIM: 'ddim',
    PLMS: 'plms',
    UNIPC: 'unipc'
} as const

export const SCHEDULER_TYPES = {
    NORMAL: 'normal',
    KARRAS: 'karras',
    EXPONENTIAL: 'exponential',
    SIMPLE: 'simple',
    DDIM_UNIFORM: 'ddim_uniform'
} as const

@Injectable()
export class WorkflowBuilder {
    private nodeCounter = 1

    // 获取新的节点 ID
    private getNextNodeId(): number {
        return this.nodeCounter++
    }

    // 获取默认采样器配置
    private getDefaultSamplerConfig(seed: number): SamplerConfig {
        return {
            name: SAMPLER_TYPES.EULER_A,
            scheduler: SCHEDULER_TYPES.KARRAS,
            steps: 30,
            cfg: 7,
            denoise: 1,
            seed
        }
    }

    // 创建采样器节点
    private createSamplerNode(config: SamplerConfig): WorkflowNode {
        return {
            id: this.getNextNodeId(),
            type: 'KSampler',
            inputs: {
                seed: config.seed,
                steps: config.steps,
                cfg: config.cfg,
                sampler_name: config.name,
                scheduler: config.scheduler,
                denoise: config.denoise,
                model: null,
                positive: null,
                negative: null,
                latent_image: null
            },
            class_type: 'KSampler'
        }
    }

    // 验证采样器配置
    private validateSamplerConfig(config: Partial<SamplerConfig>): SamplerConfig {
        const defaultConfig = this.getDefaultSamplerConfig(
            config.seed || Math.floor(Math.random() * 1000000)
        )

        return {
            name: this.validateSamplerName(config.name) || defaultConfig.name,
            scheduler: this.validateSchedulerType(config.scheduler) || defaultConfig.scheduler,
            steps: this.clamp(config.steps || defaultConfig.steps, 1, 150),
            cfg: this.clamp(config.cfg || defaultConfig.cfg, 1, 30),
            denoise: this.clamp(config.denoise || defaultConfig.denoise, 0, 1),
            seed: config.seed || defaultConfig.seed
        }
    }

    // 验证采样器名称
    private validateSamplerName(name?: string): string | undefined {
        if (!name) return undefined
        return Object.values(SAMPLER_TYPES).includes(name as any)
            ? name
            : undefined
    }

    // 验证调度器类型
    private validateSchedulerType(scheduler?: string): string | undefined {
        if (!scheduler) return undefined
        return Object.values(SCHEDULER_TYPES).includes(scheduler as any)
            ? scheduler
            : undefined
    }

    // 数值范围限制
    private clamp(value: number, min: number, max: number): number {
        return Math.min(Math.max(value, min), max)
    }

    // 构建基础的 SD 1.5 工作流
    buildSDXLWorkflow(params: {
        prompt: string
        negativePrompt: string
        width: number
        height: number
        seed: number
        sampler?: Partial<SamplerConfig>
    }): Workflow {
        const nodes: WorkflowNode[] = []
        const edges = []

        // 1. 创建 Checkpoint Loader 节点
        const checkpointNode = {
            id: this.getNextNodeId(),
            type: 'CheckpointLoaderSimple',
            inputs: {
                ckpt_name: 'sd_xl_base_1.0.safetensors'
            },
            class_type: 'CheckpointLoaderSimple'
        }
        nodes.push(checkpointNode)

        // 2. 创建正向提示词编码器
        const positivePromptNode = {
            id: this.getNextNodeId(),
            type: 'CLIPTextEncode',
            inputs: {
                text: params.prompt,
                clip: null // 将通过边连接
            },
            class_type: 'CLIPTextEncode'
        }
        nodes.push(positivePromptNode)

        // 3. 创建负向提示词编码器
        const negativePromptNode = {
            id: this.getNextNodeId(),
            type: 'CLIPTextEncode',
            inputs: {
                text: params.negativePrompt || '',
                clip: null // 将通过边连接
            },
            class_type: 'CLIPTextEncode'
        }
        nodes.push(negativePromptNode)

        // 4. 创建空潜空间
        const emptyLatentNode = {
            id: this.getNextNodeId(),
            type: 'EmptyLatentImage',
            inputs: {
                width: params.width,
                height: params.height,
                batch_size: 1
            },
            class_type: 'EmptyLatentImage'
        }
        nodes.push(emptyLatentNode)

        // 使用新的采样器配置
        const samplerConfig = this.validateSamplerConfig({
            ...params.sampler,
            seed: params.seed
        })
        const samplerNode = this.createSamplerNode(samplerConfig)
        nodes.push(samplerNode)

        // 6. 创建 VAE 解码器
        const vaeDecodeNode = {
            id: this.getNextNodeId(),
            type: 'VAEDecode',
            inputs: {
                samples: null, // 将通过边连接
                vae: null // 将通过边连接
            },
            class_type: 'VAEDecode'
        }
        nodes.push(vaeDecodeNode)

        // 7. 创建保存图片节点
        const saveImageNode = {
            id: this.getNextNodeId(),
            type: 'SaveImage',
            inputs: {
                images: null, // 将通过边连接
                filename_prefix: 'ComfyUI'
            },
            class_type: 'SaveImage'
        }
        nodes.push(saveImageNode)

        // 添加边连接
        edges.push(
            // Checkpoint 到 CLIP
            {
                from_node: checkpointNode.id,
                from_output: 'CLIP',
                to_node: positivePromptNode.id,
                to_input: 'clip'
            },
            {
                from_node: checkpointNode.id,
                from_output: 'CLIP',
                to_node: negativePromptNode.id,
                to_input: 'clip'
            },
            // Checkpoint 到 VAE
            {
                from_node: checkpointNode.id,
                from_output: 'VAE',
                to_node: vaeDecodeNode.id,
                to_input: 'vae'
            },
            // Checkpoint 到 KSampler
            {
                from_node: checkpointNode.id,
                from_output: 'MODEL',
                to_node: samplerNode.id,
                to_input: 'model'
            },
            // 提示词到 KSampler
            {
                from_node: positivePromptNode.id,
                from_output: 'CONDITIONING',
                to_node: samplerNode.id,
                to_input: 'positive'
            },
            {
                from_node: negativePromptNode.id,
                from_output: 'CONDITIONING',
                to_node: samplerNode.id,
                to_input: 'negative'
            },
            // 空潜空间到 KSampler
            {
                from_node: emptyLatentNode.id,
                from_output: 'LATENT',
                to_node: samplerNode.id,
                to_input: 'latent_image'
            },
            // KSampler 到 VAE
            {
                from_node: samplerNode.id,
                from_output: 'LATENT',
                to_node: vaeDecodeNode.id,
                to_input: 'samples'
            },
            // VAE 到保存
            {
                from_node: vaeDecodeNode.id,
                from_output: 'IMAGE',
                to_node: saveImageNode.id,
                to_input: 'images'
            }
        )

        return { nodes, edges }
    }

    // 构建 ControlNet 工作流
    buildControlNetWorkflow(params: {
        prompt: string
        negativePrompt: string
        width: number
        height: number
        seed: number
        controlImage: string
        controlType: string
        sampler?: Partial<SamplerConfig>
    }): Workflow {
        const nodes: WorkflowNode[] = []
        const edges = []

        // 1. 创建 Checkpoint Loader 节点
        const checkpointNode = {
            id: this.getNextNodeId(),
            type: 'CheckpointLoaderSimple',
            inputs: {
                ckpt_name: 'sd_xl_base_1.0.safetensors'
            },
            class_type: 'CheckpointLoaderSimple'
        }
        nodes.push(checkpointNode)

        // 2. 加载 ControlNet 模型
        const controlNetNode = {
            id: this.getNextNodeId(),
            type: 'ControlNetLoader',
            inputs: {
                control_net_name: this.getControlNetModel(params.controlType)
            },
            class_type: 'ControlNetLoader'
        }
        nodes.push(controlNetNode)

        // 3. 加载控制图片
        const loadImageNode = {
            id: this.getNextNodeId(),
            type: 'LoadImage',
            inputs: {
                image: params.controlImage
            },
            class_type: 'LoadImage'
        }
        nodes.push(loadImageNode)

        // 4. 图片预处理
        const preprocessNode = {
            id: this.getNextNodeId(),
            type: this.getPreprocessorType(params.controlType),
            inputs: {
                image: null // 将通过边连接
            },
            class_type: this.getPreprocessorType(params.controlType)
        }
        nodes.push(preprocessNode)

        // 5. 创建提示词编码器
        const positivePromptNode = {
            id: this.getNextNodeId(),
            type: 'CLIPTextEncode',
            inputs: {
                text: params.prompt,
                clip: null
            },
            class_type: 'CLIPTextEncode'
        }
        nodes.push(positivePromptNode)

        const negativePromptNode = {
            id: this.getNextNodeId(),
            type: 'CLIPTextEncode',
            inputs: {
                text: params.negativePrompt || '',
                clip: null
            },
            class_type: 'CLIPTextEncode'
        }
        nodes.push(negativePromptNode)

        // 6. 创建空潜空间
        const emptyLatentNode = {
            id: this.getNextNodeId(),
            type: 'EmptyLatentImage',
            inputs: {
                width: params.width,
                height: params.height,
                batch_size: 1
            },
            class_type: 'EmptyLatentImage'
        }
        nodes.push(emptyLatentNode)

        // 7. 创建 ControlNet 应用节点
        const applyControlNetNode = {
            id: this.getNextNodeId(),
            type: 'ControlNetApply',
            inputs: {
                conditioning: null, // positive conditioning
                control_net: null, // controlnet model
                image: null // preprocessed image
            },
            class_type: 'ControlNetApply'
        }
        nodes.push(applyControlNetNode)

        // 使用新的采样器配置
        const samplerConfig = this.validateSamplerConfig({
            ...params.sampler,
            seed: params.seed
        })
        const samplerNode = this.createSamplerNode(samplerConfig)
        nodes.push(samplerNode)

        // 9. VAE 解码器
        const vaeDecodeNode = {
            id: this.getNextNodeId(),
            type: 'VAEDecode',
            inputs: {
                samples: null,
                vae: null
            },
            class_type: 'VAEDecode'
        }
        nodes.push(vaeDecodeNode)

        // 10. 保存图片节点
        const saveImageNode = {
            id: this.getNextNodeId(),
            type: 'SaveImage',
            inputs: {
                images: null,
                filename_prefix: 'ComfyUI_ControlNet'
            },
            class_type: 'SaveImage'
        }
        nodes.push(saveImageNode)

        // 添加边连接
        edges.push(
            // 加载图片到预处理
            {
                from_node: loadImageNode.id,
                from_output: 'IMAGE',
                to_node: preprocessNode.id,
                to_input: 'image'
            },
            // Checkpoint 连接
            {
                from_node: checkpointNode.id,
                from_output: 'CLIP',
                to_node: positivePromptNode.id,
                to_input: 'clip'
            },
            {
                from_node: checkpointNode.id,
                from_output: 'CLIP',
                to_node: negativePromptNode.id,
                to_input: 'clip'
            },
            // ControlNet 应用
            {
                from_node: positivePromptNode.id,
                from_output: 'CONDITIONING',
                to_node: applyControlNetNode.id,
                to_input: 'conditioning'
            },
            {
                from_node: controlNetNode.id,
                from_output: 'CONTROL_NET',
                to_node: applyControlNetNode.id,
                to_input: 'control_net'
            },
            {
                from_node: preprocessNode.id,
                from_output: 'IMAGE',
                to_node: applyControlNetNode.id,
                to_input: 'image'
            },
            // 采样器连接
            {
                from_node: checkpointNode.id,
                from_output: 'MODEL',
                to_node: samplerNode.id,
                to_input: 'model'
            },
            {
                from_node: applyControlNetNode.id,
                from_output: 'CONDITIONING',
                to_node: samplerNode.id,
                to_input: 'positive'
            },
            {
                from_node: negativePromptNode.id,
                from_output: 'CONDITIONING',
                to_node: samplerNode.id,
                to_input: 'negative'
            },
            {
                from_node: emptyLatentNode.id,
                from_output: 'LATENT',
                to_node: samplerNode.id,
                to_input: 'latent_image'
            },
            // VAE 连接
            {
                from_node: samplerNode.id,
                from_output: 'LATENT',
                to_node: vaeDecodeNode.id,
                to_input: 'samples'
            },
            {
                from_node: checkpointNode.id,
                from_output: 'VAE',
                to_node: vaeDecodeNode.id,
                to_input: 'vae'
            },
            // 保存图片
            {
                from_node: vaeDecodeNode.id,
                from_output: 'IMAGE',
                to_node: saveImageNode.id,
                to_input: 'images'
            }
        )

        return { nodes, edges }
    }

    // 辅助方法：获取 ControlNet 模型名称
    private getControlNetModel(controlType: string): string {
        const modelMap = {
            'canny': 'control_v11p_sd15_canny.pth',
            'depth': 'control_v11f1p_sd15_depth.pth',
            'pose': 'control_v11p_sd15_openpose.pth',
            'seg': 'control_v11p_sd15_seg.pth',
            'scribble': 'control_v11p_sd15_scribble.pth',
            'normal': 'control_v11p_sd15_normal.pth'
        }
        return modelMap[controlType] || modelMap['canny']
    }

    // 辅助方法：获取预处理器类型
    private getPreprocessorType(controlType: string): string {
        const preprocessorMap = {
            'canny': 'CannyEdgePreprocessor',
            'depth': 'DepthPreprocessor',
            'pose': 'OpenposePreprocessor',
            'seg': 'SegmentationPreprocessor',
            'scribble': 'ScribblePreprocessor',
            'normal': 'NormalMapPreprocessor'
        }
        return preprocessorMap[controlType] || preprocessorMap['canny']
    }

    // 构建 LoRA 工作流
    buildLoRAWorkflow(params: {
        prompt: string
        negativePrompt: string
        width: number
        height: number
        seed: number
        loraName: string
        loraWeight: number
        sampler?: Partial<SamplerConfig>
    }): Workflow {
        const nodes: WorkflowNode[] = []
        const edges = []

        // 1. 创建 Checkpoint Loader 节点
        const checkpointNode = {
            id: this.getNextNodeId(),
            type: 'CheckpointLoaderSimple',
            inputs: {
                ckpt_name: 'sd_xl_base_1.0.safetensors'
            },
            class_type: 'CheckpointLoaderSimple'
        }
        nodes.push(checkpointNode)

        // 2. 创建 LoRA Loader 节点
        const loraNode = {
            id: this.getNextNodeId(),
            type: 'LoraLoader',
            inputs: {
                model: null, // 将通过边连接
                clip: null,  // 将通过边连接
                lora_name: params.loraName,
                strength_model: params.loraWeight,
                strength_clip: params.loraWeight
            },
            class_type: 'LoraLoader'
        }
        nodes.push(loraNode)

        // 3. 创建提示词编码器
        const positivePromptNode = {
            id: this.getNextNodeId(),
            type: 'CLIPTextEncode',
            inputs: {
                text: params.prompt,
                clip: null // 将通过边连接
            },
            class_type: 'CLIPTextEncode'
        }
        nodes.push(positivePromptNode)

        const negativePromptNode = {
            id: this.getNextNodeId(),
            type: 'CLIPTextEncode',
            inputs: {
                text: params.negativePrompt || '',
                clip: null // 将通过边连接
            },
            class_type: 'CLIPTextEncode'
        }
        nodes.push(negativePromptNode)

        // 4. 创建空潜空间
        const emptyLatentNode = {
            id: this.getNextNodeId(),
            type: 'EmptyLatentImage',
            inputs: {
                width: params.width,
                height: params.height,
                batch_size: 1
            },
            class_type: 'EmptyLatentImage'
        }
        nodes.push(emptyLatentNode)

        // 使用新的采样器配置
        const samplerConfig = this.validateSamplerConfig({
            ...params.sampler,
            seed: params.seed
        })
        const samplerNode = this.createSamplerNode(samplerConfig)
        nodes.push(samplerNode)

        // 6. 创建 VAE 解码器
        const vaeDecodeNode = {
            id: this.getNextNodeId(),
            type: 'VAEDecode',
            inputs: {
                samples: null, // 将通过边连接
                vae: null // 将通过边连接
            },
            class_type: 'VAEDecode'
        }
        nodes.push(vaeDecodeNode)

        // 7. 创建保存图片节点
        const saveImageNode = {
            id: this.getNextNodeId(),
            type: 'SaveImage',
            inputs: {
                images: null, // 将通过边连接
                filename_prefix: 'ComfyUI_LoRA'
            },
            class_type: 'SaveImage'
        }
        nodes.push(saveImageNode)

        // 添加边连接
        edges.push(
            // Checkpoint 到 LoRA
            {
                from_node: checkpointNode.id,
                from_output: 'MODEL',
                to_node: loraNode.id,
                to_input: 'model'
            },
            {
                from_node: checkpointNode.id,
                from_output: 'CLIP',
                to_node: loraNode.id,
                to_input: 'clip'
            },
            // LoRA 到 CLIP
            {
                from_node: loraNode.id,
                from_output: 'CLIP',
                to_node: positivePromptNode.id,
                to_input: 'clip'
            },
            {
                from_node: loraNode.id,
                from_output: 'CLIP',
                to_node: negativePromptNode.id,
                to_input: 'clip'
            },
            // 提示词到采样器
            {
                from_node: positivePromptNode.id,
                from_output: 'CONDITIONING',
                to_node: samplerNode.id,
                to_input: 'positive'
            },
            {
                from_node: negativePromptNode.id,
                from_output: 'CONDITIONING',
                to_node: samplerNode.id,
                to_input: 'negative'
            },
            // LoRA 模型到采样器
            {
                from_node: loraNode.id,
                from_output: 'MODEL',
                to_node: samplerNode.id,
                to_input: 'model'
            },
            // 空潜空间到采样器
            {
                from_node: emptyLatentNode.id,
                from_output: 'LATENT',
                to_node: samplerNode.id,
                to_input: 'latent_image'
            },
            // 采样器到 VAE
            {
                from_node: samplerNode.id,
                from_output: 'LATENT',
                to_node: vaeDecodeNode.id,
                to_input: 'samples'
            },
            // Checkpoint VAE 到解码器
            {
                from_node: checkpointNode.id,
                from_output: 'VAE',
                to_node: vaeDecodeNode.id,
                to_input: 'vae'
            },
            // VAE 到保存
            {
                from_node: vaeDecodeNode.id,
                from_output: 'IMAGE',
                to_node: saveImageNode.id,
                to_input: 'images'
            }
        )

        return { nodes, edges }
    }
} 